from flask import Flask, render_template, request, jsonify
import base64
import os

from deepseek_tokenizer import get_token_length

app = Flask(__name__)

# 上传文件的保存路径
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传文件夹存在
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

# 上传文件接口
@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        filename = file.filename
        token_length = get_token_length(file)
        result = {
            "token_length": token_length,
            "filename": filename,
            "message": 'File uploaded successfully'
        }
        return jsonify(result), 200

# base64文件请求接口
@app.route('/upload_base64', methods=['POST'])
def upload_base64_file():
    data = request.get_json()
    if 'file' not in data:
        return jsonify({'error': 'No file part'}), 400
    file_data = data['file']
    try:
        # 解码base64数据
        file_bytes = base64.b64decode(file_data)
        # 生成文件名
        filename = 'uploaded_file.txt'
        # 保存文件
        with open(os.path.join(app.config['UPLOAD_FOLDER'], filename), 'wb') as f:
            f.write(file_bytes)
        return jsonify({'message': 'File uploaded successfully'}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 上传文件页面
@app.route('/')
def upload_page():
    return render_template('upload.html')

if __name__ == '__main__':
    app.run(debug=True)
